import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import beta, binom

# 先验分布参数
alpha = 1
beta_param = 1

# 似然函数参数
n = 10
k = 4

# 定义先验分布
def prior(theta):
    return beta.pdf(theta, alpha, beta_param)

# 定义似然函数
def likelihood(theta):
    return binom.pmf(k, n, theta)

def posterior(theta):
    return prior(theta) * likelihood(theta)

# Metropolis-Hastings 算法参数
num_samples = 10000
theta_current = 0.5  # 初始值
samples = []

# 进行采样
for _ in range(num_samples):
    # 建议分布：对称的正态分布
    theta_proposed = np.random.normal(theta_current, 0.1)
    theta_proposed = np.clip(theta_proposed, 0, 1)  # 确保 theta 在 [0, 1] 范围内
    
    # 计算接受概率
    acceptance_ratio = posterior(theta_proposed) / posterior(theta_current)
    
    # 接受或拒绝候选样本
    if np.random.uniform(0, 1) < acceptance_ratio:
        theta_current = theta_proposed
    
    samples.append(theta_current)

    # 计算后验分布的均值和方差
posterior_samples = np.array(samples)
mean = np.mean(posterior_samples)
variance = np.var(posterior_samples)

print(f"后验分布的均值：{mean}")
print(f"后验分布的方差：{variance}")

# 绘制后验分布
plt.hist(posterior_samples, bins=30, density=True, alpha=0.6, color='g', label='MCMC samples')

# 绘制理论后验分布（Beta分布）
x = np.linspace(0, 1, 100)
theoretical_posterior = beta.pdf(x, alpha + k, beta_param + n - k)
plt.plot(x, theoretical_posterior, 'r-', label='Theoretical posterior')

plt.xlabel(r'$\theta$')
plt.ylabel('Density')
plt.title('Posterior Distribution of $\theta$')
plt.legend()
# 保存图片到本地
plt.savefig('posterior_distribution.png', dpi=300, bbox_inches='tight')